﻿//#include "stdafx.h"
#include "my_socket.h"

void socket_init()
{
#ifdef WIN32
    WSADATA m_hwsadata_;
    memset(&m_hwsadata_, 0, sizeof(m_hwsadata_));
    WSAStartup(MAKEWORD(2, 2), &m_hwsadata_);
#endif
}

// m_type暂不支持Raw Socket类型
MySocket::MySocket(MySocketType m_type):
    m_socket_(INVALID_SOCKET),m_buff_(NULL)
{
    switch(m_type)
    {
        // 声明TCP类型
        case TCP4_SOCKET:
        {
            m_socket_domain_ = AF_INET;
            m_socket_type_ = TCP4_SOCKET;
            break;
        }
        case TCP6_SOCKET:
        {
            m_socket_domain_ = AF_INET6;
            m_socket_type_ = TCP6_SOCKET;
            break;
        }
            
        // 声明TCP类型
        case UDP4_SOCKET:
        {
            m_socket_domain_ = AF_INET;
            m_socket_type_ = UDP4_SOCKET;
            break;
        }
        case UDP6_SOCKET:
        {
            m_socket_domain_ = AF_INET6;
            m_socket_type_ = UDP6_SOCKET;
            break;
        }
    }
}

MySocket::~MySocket()
{
    if (m_buff_ != NULL)
    {
        delete m_buff_;
        m_buff_ = NULL;
    }
}

MySocket& MySocket::operator=(const MySocket &cp_sock)
{
    m_socket_ = cp_sock.m_socket_;
    m_socket_type_ = cp_sock.m_socket_type_;
    m_socket_domain_ = cp_sock.m_socket_domain_;
    m_local_addr_ = cp_sock.m_local_addr_;
    m_client_addr_ = cp_sock.m_client_addr_;

    return *this;
}

int MySocket::initialize()
{
    if (m_socket_type_ == TCP4_SOCKET || m_socket_type_ == TCP6_SOCKET)
    {
        m_socket_ = socket(m_socket_domain_, SOCK_STREAM, 0);
    }
    else 
    {
        m_socket_ = socket(m_socket_domain_, SOCK_DGRAM, 0);
    }

    if (m_socket_ == INVALID_SOCKET)
    {
        return ERROR_INIT_SOCK;
    }

    // 分配接收数据缓冲区
    m_buff_ = new int8[MAX_BUFF_SIZE];

    return MYSOCKET_PROC_SUCC;
}

int MySocket::send_data(const int8 *buf, const uint32 buf_len)
{
    unsigned int send_len = 0;

    switch (m_socket_type_)
    {
        case TCP4_SOCKET:
        {
            if (is_valid_socket())
            {
                // 循环发送数据，因为TCP发包可能一次并未发完传入的数据包的
                // 大小，故而循环地发送,需改非阻塞，加入超时
                for (; send_len != buf_len;)
                {
                    int onetime_bytes = SEND(m_socket_ , buf + send_len, buf_len - send_len, 0);

                    // 发送成功，继续发。。
                    if (onetime_bytes > 0)
                    {
                        send_len += onetime_bytes;
                        continue;
                    }
                    else
                    {
                        CLOSE(m_socket_);
                        // 要么连接断开了，要么发送出错了
                        return ERROR_SEND_DATA;
                    }
                }
            }
            break;
        }
        case UDP4_SOCKET:
        {
            if (is_valid_socket())
            {
                m_sendbytes_ = SENDTO(m_socket_, 
                       buf, 
                       buf_len, 
                       0, 
                       (const struct sockaddr *)&m_client_addr_, 
                       sizeof(m_client_addr_));
            }
            break;
        }
    }

    return send_len;
}

int MySocket::recv_data(int8 **buf, const int buf_len)
{
    m_recvbytes_ = 0;
    int buf_size = buf_len < MAX_BUFF_SIZE ? buf_len : MAX_BUFF_SIZE;
    int left_len = buf_len;

    if (is_valid_socket())
    {
        switch (m_socket_type_)
        {
            case TCP4_SOCKET:
            {
                // 循环接收，需改为非阻塞，加入超时机制
                for (int recv_len = 0; recv_len != buf_len;)
                {
                    buf_size = left_len  < buf_size ? left_len : buf_size;
                    int onetime_bytes = RECV(m_socket_, m_buff_ + recv_len, buf_size, 0);
                    if (onetime_bytes > 0)
                    {
                        recv_len += onetime_bytes;
                        left_len = buf_len - recv_len;
                        continue;
                    }
                    else
                    {
                        CLOSE(m_socket_);
                        // 出错了，连接断开或其它错误
                        return ERROR_RECV_DATA;
                    }
                }

                m_recvbytes_ = buf_len;
                break;
            }
            case UDP4_SOCKET:
            {
                m_recvbytes_ = RECVFROM(m_socket_, 
                                        m_buff_, 
                                        MAX_BUFF_SIZE, 
                                        0, 
                                        (struct sockaddr *)&m_client_addr_, 
                                        sizeof(m_client_addr_));
                break;
            }
        }
    }
    *buf = m_buff_;

    return m_recvbytes_;
}

int MySocket::recv_data( int8 **buf )
{
    if (is_valid_socket())
    {
        switch (m_socket_type_)
        {
            case TCP4_SOCKET:
            {
                m_recvbytes_ = RECV(m_socket_, m_buff_, MAX_BUFF_SIZE, 0);
                *buf = m_buff_;

                if (m_recvbytes_ > 0)
                {
                    return m_recvbytes_;
                }
                else 
                {
                    CLOSE(m_socket_);
                    // 接收出错了
                    return ERROR_RECV_DATA;
                }

                break;
            }
        }
    }

    return MYSOCKET_PROC_SUCC;
}

bool MySocket::is_valid_socket()
{
    if (m_socket_ == INVALID_SOCKET)
    {
        return false;
    }

    return true;
}

int MySocket::bind_port(uint16 port)
{
    // 初始化服务器地址    
    m_local_addr_.sin_family = AF_INET;
    m_local_addr_.sin_port = htons(port);
    m_local_addr_.sin_addr.s_addr = htonl(INADDR_ANY);

    if (is_valid_socket() == false)
    {
        return ERROR_INVALID_SOCKET;
    }

    int ret = bind(m_socket_, (const struct sockaddr *)&m_local_addr_, sizeof(m_local_addr_));
    if (ret != 0)
    {
        return ERROR_BIND_PORT;
    }
    
    return MYSOCKET_PROC_SUCC;
}

int MySocket::svr_listen(const uint32 max_num)
{
    if (is_valid_socket() == false)
    {
        return ERROR_INVALID_SOCKET;
    }

    return listen(m_socket_, max_num);
}

int MySocket::accept_client(MySocket *new_socket)
{
    if (is_valid_socket() == false)
    {
        return ERROR_INVALID_SOCKET;
    }

    new_socket->alloc_buf();
    int len = sizeof(m_client_addr_);

    new_socket->m_socket_ = accept(m_socket_, (struct sockaddr *)&new_socket->m_client_addr_, &len);

    if (new_socket->m_socket_ == INVALID_SOCKET)
    {
        return ERROR_ACCEPT_CLIENT;
    }

    return MYSOCKET_PROC_SUCC;
}

int MySocket::connect_svr(const int8 *ip, uint16 port)
{
    if (is_valid_socket() == false)
    {
        return ERROR_INVALID_SOCKET;
    }

    // 初始化目标服务器IP
    m_svr_addr_.sin_family = AF_INET;
    m_svr_addr_.sin_port = htons(port);
    m_svr_addr_.sin_addr.s_addr = inet_addr(ip);

    return CONNECT(m_socket_, 
                   (const struct sockaddr *)&m_svr_addr_,
                   sizeof(m_svr_addr_));
}

int MySocket::send_file(const int8 *file_name, const int32 send_pos, int32 *send_num)
{
    int ret = MYSOCKET_PROC_SUCC;

    if (is_valid_socket())
    {
        int fh;
        if( _sopen_s( &fh, file_name, _O_RDONLY | _O_BINARY, _SH_DENYNO, 0 ) )
        {
            std::cout << "open file error:"
                      << file_name
                      << std::endl;

            return ERROR_OPENFILE;
        }

        // 定位到要发送的文件位置 
        int file_pos = _lseek(fh, send_pos, SEEK_SET);
        if (file_pos  == -1 || file_pos != send_pos)
        {
            std::cout << "seek file pos error:" << std::endl;
            return ERROR_SET_FILE_POS;
        }


        // 从定位处开始，读取文件的内容并发送
        DWORD read_num;
        int send_bytes = 0;

        // 确定读到了数据,才会进入循环
        while ((read_num = _read(fh, m_buff_ + 4, MAX_BUFF_SIZE - 4)) > 0)
        {
            // 每一帧的帧头记录数据包大小
            int net_order = htonl(read_num);
            memcpy(m_buff_, &net_order, 4);
            read_num += 4;
            int onetime_send = send_data(m_buff_, read_num);

            if (onetime_send == ERROR_SEND_DATA)
            {
                // 发送失败,记录发送结果就退出了
                *send_num = send_bytes;
                
                _close(fh);
                return ERROR_SEND_FILE;
            }

            // 统计发送数据量
            send_bytes += (onetime_send - 4);
        }

        // 发送传输结束符
        char file_ending[FILE_ENDING_LEN + 4];
        int net_order = htonl(FILE_ENDING_LEN);
        memcpy(file_ending, &net_order, sizeof(net_order));
        memcpy(file_ending + sizeof(net_order), FILE_ENDING, FILE_ENDING_LEN);

        ret = send_data(file_ending, FILE_ENDING_LEN + 4);

        // 记录一下发送了多少数据
        *send_num = send_bytes;

        _close(fh);
    }

    return ret;
}

int MySocket::recv_file( const int8 *file_name, int32 *recv_num )
{
    if (is_valid_socket())
    {
        FILE *fp;
        if (fopen_s(&fp, file_name, "ab+") != 0)
        {
            return ERROR_OPENFILE;
        }

        unsigned int onetime_bytes;
        unsigned int recv_bytes = 0;

        // 先接收帧头，看看一个文件包的大小
        while (recv_data(&m_buff_, 4)!= ERROR_RECV_DATA)
        {
            unsigned int onetime_write;
            
            int net_order = *reinterpret_cast<int *>(m_buff_);
            onetime_bytes = ntohl(net_order);

            if (recv_data(&m_buff_, onetime_bytes) == ERROR_RECV_DATA)
            {
                break;
            }

            // 检查文件传输是否结束
            if (strncmp(&m_buff_[onetime_bytes - FILE_ENDING_LEN], 
                        FILE_ENDING,
                        FILE_ENDING_LEN) == 0)
            {
                break;
            }
            
            onetime_write = static_cast<unsigned int>(fwrite(m_buff_, sizeof(char), onetime_bytes, fp));

            // 写入的大小不等于接收的大小，认为出错
            // 记录并退出
            if (onetime_bytes != onetime_write)
            {
                fclose(fp);
                *recv_num = recv_bytes;
                return ERROR_WRITE_FILE;
            }
            recv_bytes += onetime_write;    

        }

        fclose(fp);
        *recv_num = recv_bytes;
    }

    return MYSOCKET_PROC_SUCC;
}

int MySocket::alloc_buf()
{
    m_buff_ = new int8[MAX_BUFF_SIZE];

    return MYSOCKET_PROC_SUCC;
}

void MySocket::close_sock()
{
    CLOSE(m_socket_);
}

